package com.wutong.resonance.model.util;

import lombok.extern.slf4j.Slf4j;
import org.jtransforms.fft.DoubleFFT_1D;
import org.jtransforms.fft.FloatFFT_1D;

/**
 * 音频工具类
 * @author wutong
 */
@Slf4j
public class AudioUtil {

    /**
     * 计算立体声float数组指定频率振幅
     *
     * @param stereoData 立体声数据[左声道][右声道]
     * @param targetFreq 目标频率(Hz)
     * @param sampleRate 采样率(Hz)
     *
     * @return 振幅（左声道, 右声道）
     */
    public static float[] calculateFrequencyAmplitude(float[][] stereoData, float targetFreq, float sampleRate) {
        if (stereoData[0].length != stereoData[1].length) {
            throw new IllegalArgumentException("左右声道数据长度不一致");
        }

        float[] left = applyHanningWindow(stereoData[0]);
        float[] right = applyHanningWindow(stereoData[1]);

        FloatFFT_1D fft = new FloatFFT_1D(left.length);
        fft.realForward(left);
        fft.realForward(right);

        int targetBin = (int) (targetFreq * left.length / sampleRate);
        return new float[]{
                calculateMagnitude(left, targetBin),
                calculateMagnitude(right, targetBin)
        };
    }

    /**
     * 计算立体声double数组指定频率振幅
     *
     * @param stereoData 立体声数据[左声道][右声道]
     * @param targetFreq 目标频率(Hz)
     * @param sampleRate 采样率(Hz)
     *
     * @return 振幅（左声道, 右声道）
     */
    public static double[] calculateFrequencyAmplitude(double[][] stereoData, double targetFreq, double sampleRate) {
        if (stereoData[0].length != stereoData[1].length) {
            throw new IllegalArgumentException("左右声道数据长度不一致");
        }

        double[] left = applyHanningWindow(stereoData[0]);
        double[] right = applyHanningWindow(stereoData[1]);

        DoubleFFT_1D fft = new DoubleFFT_1D(left.length);
        fft.realForward(left);
        fft.realForward(right);

        int targetBin = (int) (targetFreq * left.length / sampleRate);
        return new double[]{
                calculateMagnitude(left, targetBin),
                calculateMagnitude(right, targetBin)
        };
    }

    private static float[] applyHanningWindow(float[] data) {
        float[] windowed = new float[data.length];
        for (int i = 0; i < data.length; i++) {
            float multiplier = 0.5f * (1 - (float) Math.cos(2 * Math.PI * i / (data.length - 1)));
            windowed[i] = data[i] * multiplier;
        }
        return windowed;
    }

    private static float calculateMagnitude(float[] fftData, int bin) {
        int realIndex = bin * 2;
        int imagIndex = realIndex + 1;
        float real = fftData[realIndex];
        float imag = fftData[imagIndex];
        return (float) Math.sqrt(real * real + imag * imag);
    }

    /**
     * 计算单声道float数组指定频率振幅
     * @param fftData 单声道FFT数据
     * @param bin 频率bin
     * @return 振幅
     */
    public static double calculateMagnitude(double[] fftData, int bin) {
        int realIndex = bin * 2;
        int imagIndex = realIndex + 1;
        double real = fftData[realIndex];
        double imag = fftData[imagIndex];
        return Math.sqrt(real * real + imag * imag);
    }

    /**
     * 汉明窗
     * @param data 单声道数据
     * @return 振幅
     */
    private static double[] applyHanningWindow(double[] data) {
        double[] windowed = new double[data.length];
        for (int i = 0; i < data.length; i++) {
            double multiplier = 0.5 * (1 - Math.cos(2 * Math.PI * i / (data.length - 1)));
            windowed[i] = data[i] * multiplier;
        }
        return windowed;
    }
}
